import cmd
import getopt
import sys
import sqlite3
import os
import re
import logging
'''
Created on 7 Aug 2011

@author: A Hopkins
'''

LOGGING_LEVELS= {'critical': logging.CRITICAL, 'error': logging.ERROR, 'warning':logging.WARNING, 'info':logging.INFO, 'debug':logging.DEBUG}

#def retrieveErrorDirectory():
#    # NEED EXCEPTION HANDLING HERE
#    f = open('/pentest/AH/pentest.config')
#    line = f.readline()
#    workingDir=''
#    while line:
#        if line.startswith('workingDir='):
#            m = re.search('=.*', line)
#            workingDir=m.group(0)[1:]
#        elif line.startswith('errorDir='):
#            m = re.search('=.*', line)
#            tempStr = m.group(0)[1:].replace("\"$workingDir\"", "")
#            errorDir = workingDir + tempStr
#        line = f.readline()
#    f.close()
#    return errorDir

def query(query, database, listMode):
    list = []
    if query != "":  
        try:
            logging.debug(query)
            conn = sqlite3.connect(database) 
            c = conn.cursor()
            c.execute(query)
        except sqlite3.OperationalError as strerror:
            logging.error("-- SQLITE ERROR - queryDatabase: " + str(strerror))
            sys.exit(2)  

        for row in c:
            list.append(row)
        return list
 
#*** This isn't exactly how i want it, don't like that "list mode" code in here, and "multithread code" in a different class           
def executeListModeCmd(table, cmdLine, delim):
    logging.debug ("cmdLine: " + cmdLine)
    regex2 = "\${[1-9]+}"    # look for the replacement holder, this won't match 11,12,13,14 but this shouldn't ever really matter (At the moment there aren't that many columns in the database
    regex1 = "${1}"
    # Note that you can't use the curly brackets in eclipse, so you will need to replace them, when specifying them in eclipse arguments
    # You want to check for $1 which "should" always be the IP address
    ipList = ""
    counter = 0
    for row in table:
        # at the moment, ipaddress needs to be the first column specified, would be nice if this could be detected, but maybe version 2!
        ipList += row[0]        
        if counter != len(table)-1:
            ipList += delim
        counter+=1
    
    # replace 
    for i in re.findall(regex2, cmdLine):
        logging.debug ("i: " + i)
        index = "".join(re.findall("[0-9]+", i)) 
        
        # if the index is greater than 1 (i.e. not the IP address, which at the moment, must be specified first....\${1}
        if int(index) > 1:
            cmdLine = cmdLine.replace(i,row[int(index)-1])
        # otherwise if the index is greater than 1
        elif (int(index) == 1):
            cmdLine = cmdLine.replace(i, ipList)
        else:
            logging.error("This should never be reached")

    logging.debug(cmdLine)
    os.system(cmdLine)


def usage():
    print ("=======================================================================================")
    print ("q2ex v0.1 - A Hopkins - 07 August 2011")
    print ("=======================================================================================")
    print ("-d <database>       The database you want to query against")
    print ("-q <query>          The query that you would like to make against the database")
    print ("-c <command>        The command you want to perform") 
    print ("-l                  List mode. Only one command needs to be executed, with a list of hosts,")
    print ("                    ipaddress must be first column specified, see example")
    print ("-b                  Specify the delimiter, only applies when used in list mode")
    print ("-h                  This usage")
    print ("-t <number>         Specify the number of concurrent threads")  # *** user can't currently set, but it would be trivial to implement - set as 20 in cmd.py at the moment
    print ("-T <number>         Specify the timeout value for the command. When this value is reached, the process and all children will be killed")
    print ("---------")
    print ("Examples:") # *** These are contrived examples used whilst testing, need better ones
    print ("---------")
    print ("list mode:")
    print ("python3 q2exe.py -d /job/test.db -q \'select ipaddress,port from services\' -l -c \'cat \${1}\'")
    print("")
    print ("Multithread mode:")
    print ("python3 q2exe.py -d /job/test.db -q \'select ipaddress,port from services where service = \"http\" and port = \"80\"\' -c \'ping \${1}'")
    print ("---------")
    print ("Notes...")
    print ("Use single quotes not double (apart from in where clauses)")
    print ("Eclipse specific - you cannot use curly brackets, you need to change the regex")
    print ("---------")

def main():
    
    querySpecified = False
    commandSpecified = False
    databaseSpecified = False
    listMode = False
    q=None
    delim=","
    logginglevel='warning'
    
    try:
        opts, args = getopt.getopt(sys.argv[1:], "hlc:q:d:b:t:T:o:e:n:L:O:", ["help", "listMode" , "command=", "query=", "database=", "delimiter=", "threads=", "timeout=", "outputDir=", "extension=", "logginglevel=", "outputLogFileName="])
    except getopt.GetoptError as err:
        # print help information and exit:
        usage()
        logging.error (err) # will print something like "option -a not recognized"
        sys.exit(2)
    
    for o, a in opts:
        if o in ("-h", "--help"):
            usage()
            sys.exit()
        elif o in ("-q", "--query"):
            q = a
            if q == "":
                logging.error ("No query has been specified")
                sys.exit(2)
            querySpecified = True
        elif o in ("-l", "--list"):
            listMode = True
        elif o in ("-b", "--delimiter"):
            delim = a
        elif o in ("-c", "--command"):
            command = a
            if command == "":
                logging.error ("No command has been specified")
                sys.exit(2)
            commandSpecified = True
        elif o in ("-d", "--database"):
            database = a
            if database == "":
                logging.error ("Database name has not been specified")
                sys.exit(2)
            databaseSpecified = True
        elif o in ("-t", "--threads"):
            threads = a
        elif o in ("-T", "--timeout"):
            timeout = a
        elif o in ("-o", "--outputDir"):
            outputDir = a
        elif o in ("-e", "--extension"):
            outputExtension = a
        elif o in ("-L", "--logginglevel"):
            logginglevel = a
        elif o in ("-O", "--outputLogFilename"):
            outputLog = a
                
    if commandSpecified == False:
        usage()
        logging.error ("Command has not been specified")
    elif querySpecified == False:
        usage()
        logging.error ("Query has not been specified")
    elif databaseSpecified == False:
        usage()
        logging.error ("Database has not been specified")
    elif os.path.exists(database) == False:
        logging.error ("The specified database does not exist")
    # else we are good to go
    elif commandSpecified and querySpecified and os.path.exists(database):
        logging_level=LOGGING_LEVELS.get(logginglevel, logging.NOTSET)
        logging.basicConfig(level=LOGGING_LEVELS.get('debug', logging.NOTSET), filename=outputLog, format='%(asctime)s %(levelname)s: %(message)s', datefmt='%Y-%m-%d %H:%M:%S')
        console = logging.StreamHandler()
        console.setLevel(logging_level)
        formatter = logging.Formatter('%(asctime)s %(levelname)s: %(message)s')
        console.setFormatter(formatter)
        logging.getLogger('').addHandler(console)
        
        logging.debug("================")
        logging.debug("   q2ex v0.1")
        logging.debug("================")

        result = query(q, database, listMode)
        
        # check whether we just want one command with a list of ipaddresses, or whether we want to spawn off an individual process for each result
        if listMode == False: 
            logging.debug(result)
            c = cmd.cmd(result, command, threads, timeout, outputDir, outputExtension)
        else:
            logging.debug(result)
            executeListModeCmd(result, command, delim) # timeouts are not currently implemented for this mode
            
    else:
        usage()

if __name__ == '__main__':
    main()
